import numpy as np
import tensorflow as tf
from tensorflow.keras import Sequential, layers

from dataset.captcha.captcha import load_captcha


def evaluate(net, x, t):
    y = net.predict(x)
    y = np.argmax(y, axis=-1)
    if t.ndim != 1:
        t = np.argmax(t, axis=-1)

    accuracy = np.mean(y == t)
    return accuracy


def loss_func(y_true, y_pred):
    loss_ce = tf.losses.MSE(y_true, y_pred)
    loss_ce = tf.reduce_mean(loss_ce)
    return loss_ce


# 读入数据(24,72,3)
(x_train, t_train), (x_test, t_test) = load_captcha(one_hot_label=True)
x_validation, t_validation = x_test, t_test

# 超参数
epochs = 100
batch_size = 128
learning_rate = 1e-1

train_db = tf.data.Dataset.from_tensor_slices((x_train, t_train))
train_db = train_db.batch(batch_size)

network = Sequential([
    layers.Conv2D(12, 3, 1, activation=tf.nn.leaky_relu),
    layers.MaxPooling2D(strides=2),
    layers.BatchNormalization(),
    layers.Conv2D(36, 3, 3, activation=tf.nn.leaky_relu),
    # layers.BatchNormalization(),
    # layers.Conv2D(128, (3, 5), (1, 2), activation=tf.nn.leaky_relu),
    layers.BatchNormalization(),
    layers.Flatten(),
    # layers.Dense(128 * 2),
    # layers.BatchNormalization(),
    layers.Dense(4 * 36, activation=tf.nn.sigmoid),
    layers.Reshape([4, 36])
])

network.build((None, 24, 72, 3))
network.summary()

optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
# network.compile(optimizer, loss=tf.losses.categorical_crossentropy)

for epoch in range(epochs):
    for step, (x, t) in enumerate(train_db):
        # 有的时候验证码不是这种格式，就没处理所以就不是的直接过滤
        with tf.GradientTape() as tape:
            # y
            y = network(x)

            loss_ce = loss_func(t, y)

            # 不断更新梯度
            grads = tape.gradient(loss_ce, network.trainable_variables)
            optimizer.apply_gradients(zip(grads, network.trainable_variables))

        if step % 10 == 0:
            acc = evaluate(network, x_validation, t_validation)
            print(epoch, step, "loss:", loss_ce, "acc:", acc)
    # 因为一次就已经很高了，所以直接保存模型

network.evaluate(x_test, t_test)
network.save('model.h5')
